#!/usr/bin/env python
# encoding: utf-8
"""
@author: youfeng
@email: youfeng243@163.com
@license: Apache Licence
@file: gsxt_detail_crawl.py
@time: 2018/1/3 17:46
"""
import json
import signal
import time

import click
from beanstalkc import SocketError, CommandFailed

from cache.record_cache import RecordCache
from common import util
from common.mongo import MongDb
from common.pybeanstalk import PyBeanstalk
from config.beanstalk_conf import BEANSTALK_CONF
from config.mongo_conf import MONGO_DB_SOURCE
from config.mongo_conf import MONGO_DB_TARGET
from config.redis_conf import GS_RECORD_CACHE_CONF, ANNUAL_RECORD_CACHE_CONF
from logger import Logger
from model.mq_model import MqModel

log = Logger('gsxt_detail_saver.log').get_logger()

# 当前进程运行状态
is_running = True

# 最大缓存数目
MAX_CACHE_NUM = 50


def process_quit(signo, frame):
    global is_running
    is_running = False
    log.info("收到退出进程信号...")


# 判断是否需要等待消息队列消费完成
def can_send(beanstalk, output_tube):
    MAX_BEANSTALK_NUM = 1000

    try:
        count = beanstalk.get_tube_count(output_tube)
    except CommandFailed, (_, status, results):
        if status == 'NOT_FOUND':
            log.error('获取消息队列不存在: tube = {}'.format(output_tube))
            count = 0
        else:
            log.error('获取消息队列异常: tube = {} status = {}'.format(output_tube, status))
            count = MAX_BEANSTALK_NUM
    except SocketError as e1:
        count = MAX_BEANSTALK_NUM
        beanstalk.reconnect()
        log.warn("reconnect beanstalk...")
        log.exception(e1)

    except Exception as e:
        count = MAX_BEANSTALK_NUM
        log.error("获取消息队列大小失败: ")
        log.exception(e)

    if count < MAX_BEANSTALK_NUM:
        return True

    return False


# 发送数据到消息队列
def send_parse_beanstalk(beanstalk, output_tube, detail_page_model):
    # 判断是否需要等待消息队列
    if not can_send(beanstalk, output_tube):
        return

    try:
        msg = json.dumps(detail_page_model)
        result = beanstalk.put(output_tube, msg)
        if result is None:
            log.error("数据量太大: msg len = {}".format(len(msg)))

    except SocketError as e:
        time.sleep(5)
        beanstalk.reconnect()
        log.warn("reconnect beanstalk...")
        log.exception(e)
    except Exception as e:
        log.error('捕获异常休眠:')
        log.exception(e)
        time.sleep(5)


# 存储至缓存
def store_model_to_memory(web_page_memory_cache, seed_memory_cache,
                          detail_page_model, seed_model, list_page_model,
                          detail_page_table, seed_table, list_page_table):
    # 详情页信息缓存
    if detail_page_table is not None and detail_page_model is not None:
        if detail_page_table in web_page_memory_cache:
            web_page_memory_cache[detail_page_table].append(detail_page_model)
        else:
            web_page_memory_cache[detail_page_table] = [detail_page_model]

    # 种子信息缓存
    if seed_table is not None and seed_model is not None:
        if seed_table in seed_memory_cache:
            seed_memory_cache[seed_table].append(seed_model)
        else:
            seed_memory_cache[seed_table] = [seed_model]

    # 列表页信息缓存
    if list_page_table is not None and list_page_model is not None:
        if list_page_table in web_page_memory_cache:
            web_page_memory_cache[list_page_table].append(list_page_model)
        else:
            web_page_memory_cache[list_page_table] = [list_page_model]


# 刷新缓存数据到磁盘
def flush_cache_to_db(web_page_memory_cache, seed_memory_cache,
                      detail_page_table, seed_table, list_page_table,
                      seed_db, web_page_db):
    # 详情页
    if detail_page_table is not None:
        start_time = time.time()
        detail_page_list = web_page_memory_cache.get(detail_page_table)
        if isinstance(detail_page_list, list) and len(detail_page_list) >= MAX_CACHE_NUM:
            web_page_db.insert_batch_data(detail_page_table, detail_page_list)
            log.info("定量存储详情页网页库: table = {} len = {} 耗时: {} s".format(
                detail_page_table, len(detail_page_list), time.time() - start_time))
            del web_page_memory_cache[detail_page_table][:]

    # 种子信息
    if seed_table is not None:
        start_time = time.time()
        seed_list = seed_memory_cache.get(seed_table)
        if isinstance(seed_list, list) and len(seed_list) >= MAX_CACHE_NUM:
            seed_db.insert_batch_data(seed_table, seed_list)
            log.info("定量存储种子库:  table = {} len = {} 耗时: {} s".format(
                seed_table, len(seed_list), time.time() - start_time))
            del seed_memory_cache[seed_table][:]

    # 列表页
    if list_page_table is not None:
        start_time = time.time()
        list_page_list = web_page_memory_cache.get(list_page_table)
        if isinstance(list_page_list, list) and len(list_page_list) >= MAX_CACHE_NUM:
            web_page_db.insert_batch_data(list_page_table, list_page_list)
            log.info("定量存储列表页网页库: table = {} len = {} 耗时: {} s".format(
                list_page_table, len(list_page_list), time.time() - start_time))
            del web_page_memory_cache[list_page_table][:]


# 刷新全部表到磁盘
def flush_cache_all_to_db(web_page_memory_cache, seed_memory_cache, seed_db, web_page_db):
    start_all_time = time.time()

    # 存储网页库
    for web_page_table_name, data_list in web_page_memory_cache.iteritems():
        start_time = time.time()
        web_page_db.insert_batch_data(web_page_table_name, data_list)
        log.info("定时存储网页库: table = {} len = {} 耗时: {} s".format(
            web_page_table_name, len(data_list), time.time() - start_time))
        del web_page_memory_cache[web_page_table_name][:]

    web_page_memory_cache.clear()

    # 存储种子库
    for seed_table_name, data_list in seed_memory_cache.iteritems():
        start_time = time.time()
        seed_db.insert_batch_data(seed_table_name, data_list)
        log.info("定时存储种子库:  table = {} len = {} 耗时: {} s".format(
            seed_table_name, len(data_list), time.time() - start_time))
        del seed_memory_cache[seed_table_name][:]

    seed_memory_cache.clear()

    log.info("刷新缓存数据到磁盘完成, 耗时: {} s".format(time.time() - start_all_time))


# 定时刷新缓存数据到磁盘
def flush_cache_timing(current_time, web_page_cache, seed_cache, seed_db, web_page_db):
    now_time = int(time.time())

    # 一分钟刷一次磁盘
    if now_time - current_time >= 60:
        flush_cache_all_to_db(web_page_cache, seed_cache, seed_db, web_page_db)
        return now_time

    return current_time


# 存储抓取记录到redis缓存
def store_cache_to_redis(record_cache_redis, cache_model):
    if not isinstance(cache_model, dict):
        return
    key = cache_model.get('key')
    if key is None:
        return

    value = cache_model.get('value')
    if value is None:
        return

    # 存储记录到redis
    record_cache_redis.set(key, value)


# 存储过程
def store_process(input_tube, output_tube):
    # 种子存储库
    seed_db = MongDb(MONGO_DB_SOURCE['host'], MONGO_DB_SOURCE['port'], MONGO_DB_SOURCE['db'],
                     MONGO_DB_SOURCE['username'],
                     MONGO_DB_SOURCE['password'], log=log, max_pool_size=300)

    # 网页存储库
    web_page_db = MongDb(MONGO_DB_TARGET['host'], MONGO_DB_TARGET['port'], MONGO_DB_TARGET['db'],
                         MONGO_DB_TARGET['username'],
                         MONGO_DB_TARGET['password'], log=log, max_pool_size=300)

    # 工商抓取缓存
    gs_record_cache_redis = RecordCache(GS_RECORD_CACHE_CONF['host'],
                                        GS_RECORD_CACHE_CONF['port'],
                                        GS_RECORD_CACHE_CONF['password'],
                                        GS_RECORD_CACHE_CONF['db'],
                                        GS_RECORD_CACHE_CONF['max_connections'], log=log)

    # 年报抓取缓存
    annual_record_cache_redis = RecordCache(ANNUAL_RECORD_CACHE_CONF['host'],
                                            ANNUAL_RECORD_CACHE_CONF['port'],
                                            ANNUAL_RECORD_CACHE_CONF['password'],
                                            ANNUAL_RECORD_CACHE_CONF['db'],
                                            ANNUAL_RECORD_CACHE_CONF['max_connections'], log=log)

    # 初始化监听消息队列
    beanstalk = PyBeanstalk(BEANSTALK_CONF['host'], BEANSTALK_CONF['port'])

    count = 0

    # 网页库缓存
    web_page_memory_cache = {}

    # 参数缓存
    seed_memory_cache = {}

    # 获得当前时间
    current_time = int(time.time())

    log.info("存储模块初始化完成, 监听消息队列...")
    while is_running:
        try:
            job = beanstalk.reserve(input_tube, 3)

            # 刷新数据到磁盘
            current_time = flush_cache_timing(current_time,
                                              web_page_memory_cache, seed_memory_cache,
                                              seed_db, web_page_db)

            if job is None:
                continue

            body = job.body
            job.delete()
            count += 1

            receive_data_dict = util.json_loads(body)
            if receive_data_dict is None:
                log.error('数据格式错误: msg = {}'.format(body))
                time.sleep(5)
                continue

            mq_model = MqModel(receive_data_dict)

            # 获得网页库
            detail_page_table, detail_page_model = mq_model.get_detail_page()
            if detail_page_model is not None:
                # 发送数据到在线解析队列
                send_parse_beanstalk(beanstalk, output_tube, detail_page_model)

            # 获得种子库信息
            seed_table, seed_model = mq_model.get_seed()

            # 获得列表页信息
            list_page_table, list_page_model = mq_model.get_list_page()

            # 添加数据到内存缓存
            store_model_to_memory(web_page_memory_cache, seed_memory_cache,
                                  detail_page_model, seed_model, list_page_model,
                                  detail_page_table, seed_table, list_page_table)

            # 工商 redis 缓存 可能为Null
            gs_cache_model = mq_model.get_gs_cache()

            # 存储数据到redis缓存
            store_cache_to_redis(gs_record_cache_redis, gs_cache_model)

            # 年报 redis 缓存 可能为Null
            annual_cache_model = mq_model.get_annual_cache()

            # 存储数据到redis缓存
            store_cache_to_redis(annual_record_cache_redis, annual_cache_model)

            # 存储缓存数据到mongodb
            flush_cache_to_db(web_page_memory_cache, seed_memory_cache,
                              detail_page_table, seed_table, list_page_table, seed_db, web_page_db)
        except SocketError as e:
            time.sleep(5)
            beanstalk.reconnect()
            log.warn("reconnect beanstalk...")
            log.exception(e)
        except Exception as e1:
            log.error("消费消息队列异常:")
            log.exception(e1)
            time.sleep(5)

    # 退出前先存储数据
    flush_cache_all_to_db(web_page_memory_cache, seed_memory_cache, seed_db, web_page_db)

    log.info("存储模块安全退出...")


@click.command()
@click.option('--input_tube',
              default='gsxt_store',
              help='存储消息队列')
@click.option('--output_tube',
              default='gsxt_detail_parse',
              help='解析消息队列')
def main(input_tube, output_tube):
    log.info("启动工商详情页存储进程...")
    log.info("存储消息队列: {}".format(input_tube))
    log.info("解析消息队列: {}".format(output_tube))

    # 注册优雅退出
    signal.signal(signal.SIGINT, process_quit)
    signal.signal(signal.SIGTERM, process_quit)
    signal.signal(signal.SIGQUIT, process_quit)
    signal.signal(signal.SIGUSR1, process_quit)

    # 启动存储过程
    store_process(input_tube, output_tube)


if __name__ == '__main__':
    main()
